package org.chen.handler;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.chen.model.vo.CrawlerTaskVO;
import org.chen.service.CrawlerService;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

// WebSocket处理器
@Slf4j
public class CrawlerWebSocketHandler extends TextWebSocketHandler {
    
    private static final Map<String, Set<WebSocketSession>> taskSubscriptions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionClosed(@NonNull WebSocketSession session, @NonNull CloseStatus status) throws Exception {
        log.info("WebSocket connection closed: {}, status: {}", session.getId(), status);
        super.afterConnectionClosed(session, status);
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.info("New WebSocket connection established: {}", session.getId());
    }
    
    @Override
    protected void handleTextMessage(@NonNull WebSocketSession session, @NonNull TextMessage message) {
        log.info("Received message from client: {}", message.getPayload());
        try {
            // 解析订阅消息
            JsonNode jsonNode = new ObjectMapper().readTree(message.getPayload());
            String action = jsonNode.get("action").asText();
            String taskId = jsonNode.get("taskId").asText();
            
            if ("subscribe".equals(action)) {
                subscribeToTask(taskId, session);
            } else if ("unsubscribe".equals(action)) {
                unsubscribeFromTask(taskId, session);
            }
        } catch (Exception e) {
            log.error("Error handling WebSocket message", e);
        }
    }
    
    private void subscribeToTask(String taskId, WebSocketSession session) {
        taskSubscriptions.computeIfAbsent(taskId, k -> new CopyOnWriteArraySet<>()).add(session);
    }
    
    private void unsubscribeFromTask(String taskId, WebSocketSession session) {
        Set<WebSocketSession> sessions = taskSubscriptions.get(taskId);
        if (sessions != null) {
            sessions.remove(session);
            if (sessions.isEmpty()) {
                taskSubscriptions.remove(taskId);
            }
        }
    }
    
    public void broadcastTaskUpdate(String taskId, CrawlerTaskVO taskUpdate) {
        Set<WebSocketSession> sessions = taskSubscriptions.get(taskId);
        if (sessions != null) {
            String message = null;
            try {
                message = new ObjectMapper().writeValueAsString(taskUpdate);
            } catch (JsonProcessingException e) {
                log.error("Error serializing task update", e);
                return;
            }
            
            sessions.removeIf(session -> !session.isOpen());
            
            for (WebSocketSession session : sessions) {
                try {
                    session.sendMessage(new TextMessage(message));
                } catch (IOException e) {
                    log.error("Error sending WebSocket message", e);
                }
            }
        }
    }
}